Use numbagg for `ffill` by default #8389

max-sixty · 2023-10-28T20:40:13Z

The main perf advantage here is the array doesn't need to be unstacked & stacked, which is a huge win for large multi-dimensional arrays... (I actually was hitting a memory issue running an ffill on my own, and so thought I'd get this done!)

We could move these methods to DataWithCoords, since they're almost the same implementation between a DataArray & Dataset, and exactly the same for numbagg's implementation

For transparency — the logic of "check for numbagg, check for bottleneck" I wouldn't rate at my most confident. But I'm more confident that just installing numbagg will work. And if that works well enough, we could consider only supporting numbagg for some of these in the future.

I also haven't done the benchmarks here — though the functions are relatively well benchmarked at numbagg. I'm somewhat trading off getting through these (rolling functions are coming up too) vs. doing fewer slower, and leaning towards the former, but feedback welcome...

Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

xarray/core/nputils.py

dcherian · 2023-10-29T21:52:40Z

the array doesn't need to be unstacked & stacked

Where does this happen?

We should be able to switch between bottleneck and numbagg here:

xarray/xarray/core/duck_array_ops.py

Line 691 in f63ede9

def push(array, n, axis):

and here:

xarray/xarray/core/dask_array_ops.py

Line 58 in f63ede9

def push(array, n, axis):

then you'll get dask-aware ffill for free :)

max-sixty · 2023-10-29T22:08:08Z

the array doesn't need to be unstacked & stacked

Where does this happen?

Super weird — I have a stack trace from memray that has:

da.ffill("t")
ds = self._to_temp_dataset().unstack(dim, fill_value, sparse)
result = result._unstack_once(d, stacked_indexes[d], fill_value, sparse)

It's a sampling profiler, so it misses frames. But after 10 mins of looking, I can't see how that would possibly happen — I can't see where this would get called.

Though IIRC bottleneck doesn't deal well with arrays with lots of dimensions, and that array has lots of dimensions; so I think it's still possible that something is unstacking to reduce the dimensions even I can't find it atm.

Anyway!!

We should be able to switch between bottleneck and numbagg here:

OK great!

for more information, see https://pre-commit.ci

max-sixty · 2023-11-19T00:17:31Z

I refactored this to use the duck_array_ops, that's much better. Regretfully all the previous code on this PR was just a worse version of that and was discarded.

I need to work through the version checks...

max-sixty · 2023-11-19T00:20:57Z

Unfortunately I'm seeing some segfaults when running with:

dask
target="parallel", which is the default for numbagg functions

This would not be a great experience. So I'll look at these before we consider merging; I suspect it'll be something upstream...

max-sixty · 2023-11-24T18:47:45Z

OK, this is all done I think:

Using duck array module, since

xarray/xarray/namedarray/utils.py

Line 42 in cfe4d71

def module_available(module: str) -> bool:

doesn't get the version. (Open to anyone making changes there, I don't think it's perfect, and it's now oddly named, though I probably won't be the person to make more refinements there...)
numbagg is fixed to switch to target="cpu" when running in multithreading
I also raised that issue with numba

dcherian

Just some nits

xarray/core/duck_array_ops.py

xarray/core/nputils.py

xarray/core/rolling_exp.py

dcherian · 2023-11-24T23:05:45Z

xarray/core/rolling_exp.py

            raise ImportError(
                "numbagg >= 0.2.1 is required for rolling_exp but currently numbagg is not installed"
            )
-        elif _NUMBAGG_VERSION < Version("0.2.1"):
+        elif pycompat.mod_version("numbagg") < Version("0.2.1"):


shall we boost min supported numbagg version?

I think we basically keep the 12 month cycle — it's some good infra, and I would like to be a good citizen for that process rather than deviate to save myself some boilerplate...

xarray/tests/test_missing.py

dcherian · 2023-11-24T23:07:46Z

doc/whats-new.rst

@@ -50,6 +50,10 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~

+- :py:meth:`DataArray.bfill` & :py:meth:`DataArray.ffill` now use numbagg by
+  default, which is up to 5x faster on wide arrays on multi-core machines. (:pull:`8339`)


what are "wide" arrays?

What's a better term? I agree it's not great now. A "wide" table has lots of column that it can parallelize along. i.e. the benchmarks at https://github.com/numbagg/numbagg/#nd are 5x better than bottleneck, but just above, on a single column, it's identical, because the benefit comes from the parallelism.

Co-authored-by: Deepak Cherian <[email protected]>

* upstream/main: Raise an informative error message when object array has mixed types (pydata#4700) Start renaming `dims` to `dim` (pydata#8487) Reduce redundancy between namedarray and variable tests (pydata#8405) Fix Zarr region transpose (pydata#8484) Refine rolling_exp error messages (pydata#8485) Use numbagg for `ffill` by default (pydata#8389) Fix bug for categorical pandas index with categories with EA dtype (pydata#8481) Improve "variable not found" error message (pydata#8474) Add whatsnew for pydata#8475 (pydata#8478) Allow `rank` to run on dask arrays (pydata#8475) Fix mypy tests (pydata#8476) Use concise date format when plotting (pydata#8449) Fix `map_blocks` docs' formatting (pydata#8464) Consolidate `_get_alpha` func (pydata#8465)

github-actions bot added topic-backends topic-zarr Related to zarr storage library io labels Oct 28, 2023

Illviljan reviewed Oct 29, 2023

View reviewed changes

xarray/core/nputils.py Outdated Show resolved Hide resolved

Illviljan added the run-benchmark Run the ASV benchmark workflow label Oct 29, 2023

github-actions bot added the topic-rolling label Oct 29, 2023

Use numbagg for ffill

4710af4

max-sixty force-pushed the ffill-numbagg branch from 8a222ac to 4710af4 Compare November 19, 2023 00:15

github-actions bot added topic-dask topic-arrays related to flexible array support labels Nov 19, 2023

[pre-commit.ci] auto fixes from pre-commit.com hooks

249f14c

for more information, see https://pre-commit.ci

0ed1b3c

max-sixty mentioned this pull request Nov 19, 2023

Python abort on any call of guvectorize with parallel=True with multithreading numba/numba#9288

Closed

2 tasks

max-sixty added 2 commits November 24, 2023 10:27

Merge branch 'main' into ffill-numbagg

529ec23

Use duck_array_ops for numbagg version, test import is lazy

7f519ac

max-sixty added the plan to merge Final call for comments label Nov 24, 2023

max-sixty added 2 commits November 24, 2023 10:49

650da40

f28da8f

dcherian reviewed Nov 24, 2023

View reviewed changes

max-sixty and others added 6 commits November 24, 2023 15:12

Update xarray/core/duck_array_ops.py

6c7e2bc

Co-authored-by: Deepak Cherian <[email protected]>

Update xarray/core/nputils.py

bd30ed7

Co-authored-by: Deepak Cherian <[email protected]>

Update xarray/core/rolling_exp.py

314c384

Co-authored-by: Deepak Cherian <[email protected]>

Update xarray/core/nputils.py

3973818

Co-authored-by: Deepak Cherian <[email protected]>

Merge branch 'main' into ffill-numbagg

b2a439c

b7edf31

max-sixty added 3 commits November 24, 2023 18:45

2293fb0

ddffb69

7693b85

max-sixty merged commit a683790 into pydata:main Nov 25, 2023
28 checks passed

max-sixty deleted the ffill-numbagg branch November 25, 2023 21:06

bqi343 mentioned this pull request Mar 26, 2024

Disable bottleneck by default? #7344

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use numbagg for `ffill` by default #8389

Use numbagg for `ffill` by default #8389

max-sixty commented Oct 28, 2023 •

edited

Loading

dcherian commented Oct 29, 2023

max-sixty commented Oct 29, 2023 •

edited

Loading

max-sixty commented Nov 19, 2023

max-sixty commented Nov 19, 2023

max-sixty commented Nov 24, 2023

dcherian left a comment

dcherian Nov 24, 2023

max-sixty Nov 24, 2023

dcherian Nov 24, 2023

max-sixty Nov 24, 2023

Use numbagg for ffill by default #8389

Use numbagg for ffill by default #8389

Conversation

max-sixty commented Oct 28, 2023 • edited Loading

dcherian commented Oct 29, 2023

max-sixty commented Oct 29, 2023 • edited Loading

max-sixty commented Nov 19, 2023

max-sixty commented Nov 19, 2023

max-sixty commented Nov 24, 2023

dcherian left a comment

Choose a reason for hiding this comment

dcherian Nov 24, 2023

Choose a reason for hiding this comment

max-sixty Nov 24, 2023

Choose a reason for hiding this comment

dcherian Nov 24, 2023

Choose a reason for hiding this comment

max-sixty Nov 24, 2023

Choose a reason for hiding this comment

Use numbagg for `ffill` by default #8389

Use numbagg for `ffill` by default #8389

max-sixty commented Oct 28, 2023 •

edited

Loading

max-sixty commented Oct 29, 2023 •

edited

Loading